import torch
from torch.utils.data import Dataset, DataLoader
import cv2 as cv
import os
from glob import glob
from collections import defaultdict

directory = './dataset/'


class Tobacco():

    def __init__(self, mode):
        flag = ('train_831' if mode == 'train' else 'test_351')
        self._file = os.path.join(directory, flag)
        self._name2label = {}
        self._images = defaultdict(list)
        self.X = []
        self.y = []

        for name in os.listdir(self._file):
            self._name2label[name] = len(self._name2label.keys())
            self._images[name] = glob(os.path.join(self._file, name, '*.jpg'))

        for label, images in self._images.items():
            for img in images:
                self.y.append(self._name2label[label])
                self.X.append(cv.imread(img))

    def __getitem__(self, item):
        raise NotImplementedError

    def __len__(self):
        raise NotImplementedError


def file():
    file = './dataset/train_831/B2F/00001.jpg'
    img = cv.imread(file)
    return


if __name__ == '__main__':
    data = Tobacco('train')

    print()
